#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/3/6 16:12
# @Author  : wanghaoran
# @File    : keywords_retrieval.py
from advanced_module.base import BaseRetrievalEngineTool
import requests
from elasticsearch_dsl import Search
from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk


class KeywordRetrieval(BaseRetrievalEngineTool):
    def __init__(self, es_host, es_port):
        self.es_host = es_host
        self.es_port = es_port

    def update_info(self, ids, fields, values):
        self.es = Elasticsearch(f"http://{self.es_host}:{self.es_port}", connections_per_node=15, timeout=100, max_retries=10,
                                retry_on_timeout=True)
        doc = dict()
        for k, v in zip(fields, values):
            doc[k] = v

        actions = [
            {
                "_op_type": "update",
                "_index": "detail_index",
                "_id": doc_id,
                "doc": doc
            }
            for doc_id in ids
        ]
        bulk(self.es, actions)

    def get_results(self, query, base_name_list, topk, query_vector=None, embedding_model=None, similarity_type=None):
        # print("base_name_list", base_name_list)
        s = Search().query('multi_match', query=query, fields=['text', 'detail_text'])
        s = s.query('terms', base__name=base_name_list)

        es_res = s.execute()
        es_results = [hit for hit in es_res.hits.hits]
        keyword_results = []
        for r in es_results[:topk]:
            keyword_results.append({
                "uid": r["_source"]["uid"],
                "text": r["_source"]["text"],
                "detail_text": r["_source"]["detail_text"],
                "source": r["_source"]["source"]["name"] if "name" in r["_source"]["source"] else "无",
                "base__name": r["_source"]["base"]["name"],
                "distance": r["_score"],
                "comment": "keyword"
            })
        return keyword_results


class WebQueryEngine(BaseRetrievalEngineTool):
    def __init__(self, url, key, language):
        self.url = url
        self.key = key
        self.mkt = language

    def get_results(self, query):
        params = {"q": query, "mkt": self.mkt}
        headers = {"Ocp-Apim-Subscription-Key": self.key}

        # Call the API
        try:
            response = requests.get(self.url, headers=headers, params=params)
            response.raise_for_status()
            res = response.json()["webPages"]["value"]
            info = []
            for r in res:
                info.append([r["name"], r["snippet"]])

            return info
        except Exception as ex:
            print(ex)
            return []


# class VectorRetrieval(BaseRetrievalEngineTool):
#     def __init__(self):
#         pass
#
#     def get_results(self, query, base_name_list, topk, query_vector, embedding_model, similarity_type):
#         if similarity_type == 'L2':
#             res = list(embedding_model.objects.filter(base__name__in=base_name_list, status=1).order_by(L2Distance('embedding', query_vector))[:topk].values('embedding','text__text','text__detail_text', 'text__parent_id', 'text__meta_data', 'text__source__name', 'text__uid','base__name', 'base__project__name'))
#             for i in range(len(res)):
#                 res[i]['distance'] = calculate_distance(query_vector, res[i]['embedding'], 'L2')
#                 res[i].pop('embedding')
#
#         else:
#             res = list(embedding_model.objects.filter(base__name__in=base_name_list, status=1).order_by(
#                 -MaxInnerProduct('embedding', query_vector))[:topk].values('embedding', 'text__text', 'text__detail_text',
#                                                                      'text__parent_id', 'text__meta_data',
#                                                                      'text__source__name', 'text__uid', 'base__name',
#                                                                      'base__project__name'))
#             for i in range(len(res)):
#                 res[i]['distance'] = calculate_distance(query_vector, res[i]['embedding'], 'IP')
#                 res[i].pop('embedding')
#
#         vector_results = []
#         for r in res[:topk]:
#             vector_results.append({
#                 "uid": r["text__uid"],
#                 "question": r["text__text"],
#                 "detail_text": r["text__detail_text"],
#                 "score": r["distance"],
#                 "base": r["base__name"],
#                 "comment": "vector"
#             })
#         print(vector_results)
#         return vector_results